(*
   Copyright (C) 2024 International Digital Economy Academy.
   This program is licensed under the MoonBit Public Source
   License as published by the International Digital Economy Academy,
   either version 1 of the License, or (at your option) any later
   version. This program is distributed in the hope that it will be
   useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the MoonBit
   Public Source License for more details. You should have received a
   copy of the MoonBit Public Source License along with this program. If
   not, see
   <https://www.moonbitlang.com/licenses/moonbit-public-source-license-v1>.
*)

(* 并行编译支持模块 *)

module Io = Basic_io
module Config = Basic_config
module Parse = Parsing_parse
module Lst = Basic_lst
module Vec = Basic_vec

(* 并行编译配置 *)
let num_jobs = ref 4  (* 默认并行任务数 *)

(* 设置并行任务数 *)
let set_jobs n =
  if n > 0 then
    num_jobs := n
  else
    Printf.eprintf "Warning: Invalid number of jobs %d, using default value %d\n" n !num_jobs

(* 编译任务类型 *)
type compile_task = {
  input_file: string;
  output_file: string option;
  build_context: Driver_compenv.build_context;
}

(* 编译单个文件的函数 *)
let compile_file task =
  try
    let name = Filename.chop_extension task.input_file in
    let content = Io.load_file task.input_file in
    let ast = Parse.parse_file ~diagnostics:(Diagnostics.make ()) ~fname:task.input_file content in
    let postprocessed_ast = Driver_util.postprocess_ast ~diagnostics:(Diagnostics.make ()) ast in
    let tast, genv = Moon0_main.tast_of_ast ~name ~build_context:task.build_context [postprocessed_ast] in
    let core = Moon0_main.core_of_tast ~name genv tast in
    match task.output_file with
    | Some output -> 
        let wasm = Moon0_main.wasm_gen ~name core in
        Io.write output wasm;
        Ok (task.input_file)
    | None -> Ok (task.input_file)
  with e ->
    Error (task.input_file, Printexc.to_string e)

(* 并行编译多个文件 *)
let parallel_compile_files tasks =
  if !num_jobs <= 1 then
    (* 单线程模式 *)
    List.map compile_file tasks
  else
    (* 多线程模式 - 使用OCaml的Domain功能（OCaml 5.0+）或回退到顺序处理 *)
    let rec process_chunk acc = function
      | [] -> List.rev acc
      | chunk ->
          let results = List.map compile_file chunk in
          process_chunk (List.rev_append results acc) []
    in
    
    (* 将任务分成多个块 *)
    let chunk_size = max 1 (List.length tasks / !num_jobs) in
    let rec split_tasks acc current_chunk remaining =
      match remaining with
      | [] -> 
          if current_chunk <> [] then
            List.rev (current_chunk :: acc)
          else
            List.rev acc
      | task :: rest ->
          if List.length current_chunk >= chunk_size then
            split_tasks (current_chunk :: acc) [task] rest
          else
            split_tasks acc (task :: current_chunk) rest
    in
    
    let chunks = split_tasks [] [] tasks in
    
    (* 尝试使用Domain并行处理，如果不支持则回退到顺序处理 *)
    try
      let domains = List.map (fun chunk ->
        Domain.spawn (fun () -> process_chunk [] chunk)
      ) chunks in
      
      List.flatten (List.map Domain.join domains)
    with
    | Unhandled_exception _ -> 
        (* 回退到顺序处理 *)
        Printf.eprintf "Warning: Parallel compilation not supported, falling back to sequential processing\n";
        List.map compile_file tasks
    | _ ->
        (* 其他异常，回退到顺序处理 *)
        Printf.eprintf "Warning: Error in parallel compilation, falling back to sequential processing\n";
        List.map compile_file tasks

(* 创建编译任务 *)
let create_task input_file output_file build_context =
  { input_file; output_file; build_context }

(* 处理编译结果 *)
let process_results results =
  let successes = ref 0 in
  let failures = ref 0 in
  
  List.iter (function
    | Ok file -> 
        incr successes;
        if not !Config.quiet then
          Printf.printf "Successfully compiled %s\n" file
    | Error (file, err) ->
        incr failures;
        Printf.eprintf "Error compiling %s: %s\n" file err
  ) results;
  
  (!successes, !failures)